﻿#include <algorithm>
#include <gtest/gtest.h>
#include <iostream>
#include <learn/math/graph.h>
#include <learn/stl/tree.h>
#include <learn/utils/output.h>
#include <vector>

using namespace xi;
using namespace xi::math;

TEST(TestGraph, sparse)
{
    Graph_Sparse g(6);
    std::vector<std::size_t> vs(g.size());
    std::iota(vs.begin(), vs.end(), 0);

    g.create_directed_edges(vs[0], {vs[1], vs[3]}, {2, 3});
    g.create_directed_edges(vs[1], {vs[4]}, {1});
    g.create_directed_edges(vs[1], {vs[4]}, {1});
    g.create_directed_edges(vs[1], {vs[4]}, {1});
    g.create_directed_edges(vs[2], {vs[4], vs[5]}, {1, 1});
    g.create_directed_edges(vs[2], {vs[4], vs[5]}, {1, 1});
    g.create_directed_edges(vs[2], {vs[4], vs[5]}, {1, 1});
    g.create_directed_edges(vs[3], {vs[1]}, {1});
    g.create_directed_edges(vs[4], {vs[3]}, {1});
    g.create_directed_edges(vs[5], {vs[5]}, {1});

    std::cout << g << std::endl;
    g.remove_duplicates();
    std::cout << g << std::endl;
    g.transpose();
    std::cout << g << std::endl;
}

TEST(TestGraph, dense)
{
    Graph_Dense g(6);
    std::vector<std::size_t> vs(g.size());
    std::iota(vs.begin(), vs.end(), 0);

    g.create_directed_edges(vs[0], {vs[1], vs[3]}, {2, 3});
    g.create_directed_edges(vs[1], {vs[4]}, {1});
    g.create_directed_edges(vs[1], {vs[4]}, {1});
    g.create_directed_edges(vs[1], {vs[4]}, {1});
    g.create_directed_edges(vs[2], {vs[4], vs[5]}, {1, 1});
    g.create_directed_edges(vs[2], {vs[4], vs[5]}, {1, 1});
    g.create_directed_edges(vs[2], {vs[4], vs[5]}, {1, 1});
    g.create_directed_edges(vs[3], {vs[1]}, {1});
    g.create_directed_edges(vs[4], {vs[3]}, {1});
    g.create_directed_edges(vs[5], {vs[5]}, {1});

    std::cout << g << std::endl;
    g.remove_duplicates();
    std::cout << g << std::endl;
    std::cout << g.shortest_path(vs[0], vs[4], 10) << std::endl;
    g.transpose();
    std::cout << g << std::endl;
}

TEST(TestGraph, bfs)
{
    Graph_Sparse g(8);
    std::vector<std::size_t> vs(g.size());
    std::iota(vs.begin(), vs.end(), 0);

    g.create_undirected_edges(vs[0], {vs[1]}, {1});
    g.create_undirected_edges(vs[1], {vs[2]}, {1});
    g.create_undirected_edges(vs[2], {vs[3]}, {1});
    g.create_undirected_edges(vs[3], {vs[4], vs[5]}, {1, 1});
    g.create_undirected_edges(vs[4], {vs[5], vs[6], vs[7]}, {1, 1, 1});
    g.create_undirected_edges(vs[5], {vs[6]}, {1});
    g.create_undirected_edges(vs[6], {vs[7]}, {1});

    auto tree = g.bfs(vs[2]);

    // 前序遍历获得节点的距离
    std::vector<float> dist;
    auto out = [&dist](std::decay_t<decltype(*tree)>::Node *node) { dist.push_back(*node->value); };
    tree->preorder_walk(tree->root(), out);

    std::vector<float> expected_dist = {0, 1, 2, 3, 2, 3, 1, 2};
    EXPECT_EQ(dist, expected_dist);
}

TEST(TestGraph, dfs)
{
    Graph_Sparse g(6);
    std::vector<std::size_t> vs(g.size());
    std::iota(vs.begin(), vs.end(), 0);

    g.create_undirected_edges(vs[0], {vs[1], vs[2], vs[3]}, {1, 1, 1});
    g.create_undirected_edges(vs[1], {vs[2]}, {1});
    g.create_undirected_edges(vs[2], {vs[3]}, {1});
    g.create_undirected_edges(vs[3], {vs[4]}, {1});
    g.create_undirected_edges(vs[4], {vs[5]}, {1});

    auto tree = g.dfs(vs[2]);
    std::cout << g << std::endl;

    // 前序遍历获得节点的距离
    std::cout << tree << std::endl;

    // std::vector<float> expected_dist = {0, 1, 2, 3, 2, 3, 1, 2};
    // EXPECT_EQ(dist, expected_dist);
}